package com.zw.utils;

import java.math.BigInteger;
import java.security.SecureRandom;
import java.util.Base64;
import java.util.List;

public class PaillierUtil {

    private static final SecureRandom SECURE_RANDOM = new SecureRandom();

    //---------------- 密钥生成 -------------------
    public static PaillierKey.PublicKey generatePublicKey(BigInteger n, BigInteger g) {
        return new PaillierKey.PublicKey(n, g);
    }

    public static PaillierKey.PrivateKey generatePrivateKey(BigInteger lambda, BigInteger mu, BigInteger n) {
        return new PaillierKey.PrivateKey(lambda, mu, n);
    }

    /**
     * 生成完整的密钥对
     * @param bitLength 素数位数（推荐>=1024）
     */
    public static PaillierKeyPair generateKeyPair(int bitLength) {
        BigInteger p = BigInteger.probablePrime(bitLength, SECURE_RANDOM);
        BigInteger q = BigInteger.probablePrime(bitLength, SECURE_RANDOM);
        BigInteger n = p.multiply(q);
        BigInteger g = n.add(BigInteger.ONE);
        BigInteger lambda = lcm(p.subtract(BigInteger.ONE), q.subtract(BigInteger.ONE));
        BigInteger mu = L(g.modPow(lambda, n.pow(2)), n).modInverse(n);

        return new PaillierKeyPair(
                new PaillierKey.PublicKey(n, g),
                new PaillierKey.PrivateKey(lambda, mu, n)
        );
    }
    //---------------- 加密 -------------------
    /**
     * 加密整数（输出Base64字符串）
     */
    public static String encrypt(BigInteger plaintext, PaillierKey.PublicKey publicKey) {
        validatePlaintext(plaintext, publicKey.getN());

        BigInteger r = generateRandomR(publicKey.getN());
        BigInteger ciphertext = publicKey.getG().modPow(plaintext, publicKey.getNSquared())
                .multiply(r.modPow(publicKey.getN(), publicKey.getNSquared()))
                .mod(publicKey.getNSquared());

        return Base64.getEncoder().encodeToString(ciphertext.toByteArray());
    }

    //---------------- 解密 -------------------
    /**
     * 解密Base64密文字符串
     */
    public static BigInteger decrypt(String ciphertextBase64, PaillierKey.PrivateKey privateKey) {
        byte[] cipherBytes = Base64.getDecoder().decode(ciphertextBase64);
        BigInteger ciphertext = new BigInteger(cipherBytes);

        BigInteger nSquared = privateKey.getN().pow(2);
        return L(ciphertext.modPow(privateKey.getLambda(), nSquared), privateKey.getN())
                .multiply(privateKey.getMu())
                .mod(privateKey.getN());
    }

    //---------------- 辅助方法 -------------------
    private static BigInteger generateRandomR(BigInteger n) {
        BigInteger r;
        do {
            r = new BigInteger(n.bitLength(), SECURE_RANDOM);
        } while (r.compareTo(n) >= 0 || !r.gcd(n).equals(BigInteger.ONE));
        return r;
    }

    private static void validatePlaintext(BigInteger m, BigInteger n) {
        if (m.compareTo(BigInteger.ZERO) < 0 || m.compareTo(n) >= 0) {
            throw new IllegalArgumentException("明文必须满足 0 ≤ m < n");
        }
    }

    private static BigInteger L(BigInteger x, BigInteger n) {
        return x.subtract(BigInteger.ONE).divide(n);
    }

    private static BigInteger lcm(BigInteger a, BigInteger b) {
        return a.multiply(b).divide(a.gcd(b));
    }

    // 密钥对包装类
    public static class PaillierKeyPair {
        private final PaillierKey.PublicKey publicKey;
        private final PaillierKey.PrivateKey privateKey;

        public PaillierKeyPair(PaillierKey.PublicKey publicKey, PaillierKey.PrivateKey privateKey) {
            this.publicKey = publicKey;
            this.privateKey = privateKey;
        }

        // Getters
        public PaillierKey.PublicKey getPublicKey() { return publicKey; }
        public PaillierKey.PrivateKey getPrivateKey() { return privateKey; }
    }

    //---------------- 同态加法 -------------------
    /**
     * 密文加法（输入输出均为Base64字符串）
     */
    public static String add(String ciphertext1Base64, String ciphertext2Base64, PaillierKey.PublicKey publicKey) {
        BigInteger c1 = decodeBase64(ciphertext1Base64);
        BigInteger c2 = decodeBase64(ciphertext2Base64);
        BigInteger result = c1.multiply(c2).mod(publicKey.getNSquared());
        return encodeBase64(result);
    }
    /**
     * 对密文集合进行同态累加（输入输出均为Base64字符串）
     * @param ciphertextsBase64 密文集合（直接从数据库读取的字段值）
     * @param publicKey 公钥
     * @return 累加结果的Base64密文
     */
    public static String batchAdd(List<String> ciphertextsBase64, PaillierKey.PublicKey publicKey) {
        if (ciphertextsBase64 == null || ciphertextsBase64.isEmpty()) {
            return encrypt(BigInteger.ZERO, publicKey); // 返回加密的0
        }

        // 初始化累加器为第一个密文
        BigInteger accumulator = base64ToBigInteger(ciphertextsBase64.get(0));
        BigInteger nSquared = publicKey.getNSquared();

        // 从第二个元素开始遍历累加
        for (int i = 1; i < ciphertextsBase64.size(); i++) {
            BigInteger currentCipher = base64ToBigInteger(ciphertextsBase64.get(i));
            accumulator = accumulator.multiply(currentCipher).mod(nSquared);
        }

        return Base64.getEncoder().encodeToString(accumulator.toByteArray());
    }

    // Base64转BigInteger（复用已有方法）
    private static BigInteger base64ToBigInteger(String base64) {
        return new BigInteger(Base64.getDecoder().decode(base64));
    }

    private static BigInteger decodeBase64(String base64Str) {
        byte[] bytes = Base64.getDecoder().decode(base64Str);
        return new BigInteger(bytes);
    }

    private static String encodeBase64(BigInteger number) {
        return Base64.getEncoder().encodeToString(number.toByteArray());
    }
    public static void main(String[] args) {
        // 生成密钥对（测试使用512位，生产至少2048位）
        PaillierUtil.PaillierKeyPair keyPair = PaillierUtil.generateKeyPair(512);

        // 提取公钥参数
        BigInteger n = keyPair.getPublicKey().getN();
        BigInteger g = keyPair.getPublicKey().getG();

        // 提取私钥参数
        BigInteger lambda = keyPair.getPrivateKey().getLambda();
        BigInteger mu = keyPair.getPrivateKey().getMu();
        BigInteger privateKeyN = keyPair.getPrivateKey().getN();

        // 打印配置参数（用于application.yml）
        System.out.println("=== 公钥配置 ===");
        System.out.println("paillier.publicKey.n: " + n.toString());
        System.out.println("paillier.publicKey.g: " + g.toString());

        System.out.println("\n=== 私钥配置 ===");
        System.out.println("paillier.privateKey.lambda: " + lambda.toString());
        System.out.println("paillier.privateKey.mu: " + mu.toString());
        System.out.println("paillier.privateKey.n: " + privateKeyN.toString());
    }
}
